
#include <assert.h>
#include <autoconf.h>
#include <libut.h>
#include <nfs/nfs.h>
#include <string.h>

#include "../macros.h"
#include "../mapping.h"
#include "../swaptable.h"
#include "../task.h"
#include "../ut_manager/ut.h"
#include "../vmem_layout.h"

#include "../frametable.h"

#include "../debug.h"
#define verbose SOS_DEBUG_SWAP
#include <sys/debug.h>
#include <sys/panic.h>

/* Metadata about swapped out frames. Used to restore */
typedef struct swapped_frames {
    seL4_Word swap_num;     /* Swap number for frame */
    frame_mapping mapping;  /* Copy of frame's mapping data */
    /* Handle race conditions on swapped frames */
    seL4_Word sos_vaddr;    /* SOS window address for wherever the frame was */
    int abort;              /* This gets set if process faults during swapout */
    int inflight;           /* Set to 1 while being swapped, 0 when done */
    frame_swapin_cb cb;     /* Callback for an aborted frame */
    void *cookie;           /* Cookie for an aborted frame's callback */
    /* Treat swap metadata as a hash table */
    UT_hash_handle hh;      /* libut treats this type as a hash table */
} swapped_frames;

/* Callbacks passed into frame_swapout. */
typedef void (*frame_swapout_cb)(int err, seL4_Word frame_number,
                                 seL4_Word swap_number, void *cookie);

/* Data used by swapout callbacks */
typedef struct frame_swapout_data {
    frame_table_entry *fte;
    seL4_Word swap_num;
    swapped_frames *meta;
    /* Loop info */
    size_t sent;
    /* Info to call back */
    frame_swapout_cb cb;
    void *cookie;
} frame_swapout_data;

/* Data used by swapin callbacks */
typedef struct frame_swapin_data {
    seL4_Word swap_number;
    seL4_Word frame_number;
    seL4_Word sos_vaddr;
    /* Loop variables */
    size_t sent;
    /* Info to reply to caller */
    swapped_frames *meta;
    frame_swapin_cb cb;
    void *cookie;
} frame_swapin_data;

extern fhandle_t mnt_point;
frame_table_entry *_swap_queue;
static bool _swap_file_initialised;
static fhandle_t _swap_file;
static swapped_frames *_swap_meta = SWAP_META_START;

extern seL4_Word _cur_frames;

void _frame_chance(frame_table_entry *fte);
static void _swap_file_ready(uintptr_t token, enum nfs_stat status,
                             fhandle_t *fh, fattr_t *fattr);
void _swap_file_init(task_cb cb, void *cookie);

int frame_swapout(frame_swapout_cb cb, void *cookie);
static int _frame_swapout(frame_swapout_data *data);
static int _do_swapout(void *cookie);
static void _finalise_swapout(uintptr_t token, enum nfs_stat status,
                              fattr_t *fattr, int count);

static int _frame_swapin(seL4_Word swap_num, seL4_Word sos_vaddr);
static void _swapin_gotframe(seL4_Word sos_vaddr, void *cookie);
static void _finalise_swapin(uintptr_t token, enum nfs_stat status,
                             fattr_t *fattr, int count, void *frame);

typedef struct swap_file_init_data {
    task_cb cb;
    void *cookie;
} swap_file_init_data;

/* Initialise the fhandle_t for the swap file */
void _swap_file_init(task_cb cb, void *cookie) {
    enum rpc_stat err;

    // Put args on the heap for later
    swap_file_init_data *data = malloc(sizeof(struct swap_file_init_data));
    conditional_panic(!data, "Kernel OOM - cannot initialise swap file");
    data->cb = cb;
    data->cookie = cookie;

    // Build swap metadata table
    assert(SWAP_META_START > DMA_VEND);
    assert(SOS_SWAP_MAX * sizeof(struct swapped_frames) <= SWAP_META_SIZE);
    err = _range_alloc_retype_and_map(SWAP_META_START, SWAP_META_SIZE);
    conditional_panic(err, "Failed to allocate swap metadata table");

    err = nfs_lookup(&mnt_point, SWAP_FILE_NAME,
                     &_swap_file_ready, (uintptr_t) data);
    if (err) {
        dprintf(1, "swap_file_init: failed to create file. NFS rpc_stat = %d\n",
                err);
        assert(!"swap_file_init: failed to create");
    }
    return;
}

/* Callback for nfs_create / nfs_lookup */
static void _swap_file_ready(uintptr_t token, enum nfs_stat status,
                             fhandle_t *fh, fattr_t *fattr) {
    int err;
    sattr_t sattr;
    swap_file_init_data *data = (swap_file_init_data *) token;

    if (status == NFSERR_NOENT) {
        dprintf(1, "Swap file does not exist. Creating...\n");

        sattr = (sattr_t) {
            .mode = 0b110110110, // RW
            .uid = 1,
            .gid = 1,
            .size = 0,
            .atime = { 0 },
            .mtime = { 0 }
        };

        err = nfs_create(&mnt_point, SWAP_FILE_NAME, &sattr,
                         &_swap_file_ready, 0);
        if (err) {
            dprintf(1, "Could not create swap file. (%d)\n", err);
            _swap_file_initialised = false;
        }
    } else if (status) {
        // Something else is wrong - disable swap system-wide.
        dprintf(1, "Unknown error - could not initialise swap file.\n");
        _swap_file_initialised = false;
    } else {
        // XXX Check we can read and write to the file...
        // Copy the filehandle and proceed
        dprintf(1, "Swap file initialised successfully.\n");
        _swap_file_initialised = true;
        memcpy(&_swap_file, (void *) fh, sizeof(struct fhandle));
        // Perform callback
        if (data->cb) (*(data->cb))(data->cookie);
    }
    free(data);
    return;
}

/* Perform all swap-related bookkeeping, and set up the frame eviction. */
/* This is called in a synchronous context, and it is the right spot
   to inform processes that their pages have been swapped out. */
static int _frame_swapout(frame_swapout_data *data) {
    int err;
    seL4_Word swap_num;
    swapped_frames *meta;
    seL4_Word sos_vaddr;

    frame_table_entry *fte = data->fte;

    dprintf(5, "frame_swapout: evicting FTE %p (frame %d)\n",
            fte, _fte_to_framenumber(fte));

    // Get a free swap number
    err = sos_swap_alloc(&swap_num);
    if (err) return err;

    // Point to metadata array slot
    meta = _swap_meta + swap_num;
    bzero((void *) meta, sizeof(struct swapped_frames));

    // Copy the relevant stuff up into meta and data
    meta->abort = 0;
    meta->inflight = 1;
    meta->swap_num = swap_num;
    memcpy(&(meta->mapping), &(fte->mapping), sizeof(struct frame_mapping));
    data->swap_num = swap_num;
    data->meta = meta;
    meta->sos_vaddr = frame_translate_number(_fte_to_framenumber(fte));

    dprintf(8, "frame_swapout: mapping info pcb %p\tproc_vaddr %pswappable %d\n",
            fte->mapping.pcb, fte->mapping.proc_vaddr, fte->swappable);
    assert(fte->mapping.pcb);

    // Ungrant it from the given process.
    // N.B. frame_ungrant will make sure it is unmapped, if it isn't already so.
    sos_vaddr = frame_translate_number(_fte_to_framenumber(fte));
    dprintf(8, "frame_swapout: fte %p\tsos_vaddr %p\tframe_number %d\n",
            fte, sos_vaddr, _fte_to_framenumber(fte));
    err = frame_ungrant(sos_vaddr, fte->mapping.pcb->cspace, fte->mapping.proc_cap);
    if (err) return err;

    // Update the relevant process page table.
    // sos_page_swapout is sufficient.
    err = sos_page_swapout(meta->mapping.pcb, meta->mapping.proc_vaddr, swap_num);
    if (err) return err;

    // Copy FTE into heap for our coming callbacks
    data->fte = fte;

    dprintf(6, "frame_swapout: synchronous component complete\n");

    // One of two continuations - swap file is only created when needed.
    if (_swap_file_initialised) {
        // Set up eviction
        return _do_swapout((void *) data);
    } else {
        // Initialise swap file, then continue with eviction
        _swap_file_init(&_do_swapout, data);
        return FRAME_TABLE_SUCCESS;
    }
}

/* Task triggered after swap file initialised */
static int _do_swapout(void *cookie) {
    int err;

    // Actually set up the eviction
    frame_swapout_data *data = (frame_swapout_data *) cookie;

    // Check for abort - process could have faulted while we created file
    if (data->meta->abort) {
        dprintf(6, "frame_swapout: aborting operation (lost race)\n");
        (*(data->cb))(FRAME_TABLE_SWAPRACE, _fte_to_framenumber(data->fte),
                      data->swap_num, data->cookie);
        _frame_swapin(data->swap_num, data->meta->sos_vaddr);
        free(data);
        return;
    }

    // Need the aligned sos_vaddr window address for fte
    seL4_Word sos_vaddr = _fte_to_framenumber(data->fte);
    sos_vaddr = PAGE_ALIGN(frame_translate_number(sos_vaddr));

    dprintf(6, "frame_do_swapout: writing to %p\n", sos_vaddr);

    int offset = (data->swap_num) * PAGE_SIZE;

    // Initiate write operation, hope it all works in one go
    err = (int) nfs_write(&_swap_file, offset, PAGE_SIZE, sos_vaddr,
                          &_finalise_swapout, (uintptr_t) data);

    return SOS_TASK_SUCCESS;
}

/* Task triggered after frame_swapout is complete */
static void _finalise_swapout(uintptr_t token, enum nfs_stat status,
                              fattr_t *fattr, int count) {
    int err;
    frame_swapout_data *data = (frame_swapout_data *) token;

    /* We might need to do some more writing. */
    int offset;
    int remain;
    seL4_Word sos_vaddr;

    dprintf(6, "frame_swapout_finalise: nfs_write status is %d\n", status);

    // Check if operation has to be aborted (process faulted?)
    if (data->meta->abort) {
        dprintf(6, "frame_swapout: aborting operation (lost race)\n");

        (*(data->cb))(FRAME_TABLE_SWAPRACE, _fte_to_framenumber(data->fte),
                      data->swap_num, data->cookie);
        // Redo mappings, free up metadata, wake up faulted process
        _frame_swapin(data->swap_num, data->meta->sos_vaddr);

        free(data);
        return;
    }

    data->sent += count;
    // If there's still more of the page to write, write more
    if (data->sent != PAGE_SIZE) {
        offset = ((data->swap_num) * PAGE_SIZE) + data->sent;
        remain = PAGE_SIZE - data->sent;

        sos_vaddr = _fte_to_framenumber(data->fte);
        sos_vaddr = PAGE_ALIGN(frame_translate_number(sos_vaddr));
        sos_vaddr += data->sent;

        dprintf(6, "frame_swapout_finalise: %dB from offset %d into %p\n",
                remain, offset, sos_vaddr);

        err = (int) nfs_write(&_swap_file, offset, remain, sos_vaddr,
                              &_finalise_swapout, (uintptr_t) data);
        return;
    }

    // XXX M7 there will be a queue of operations hanging out on this swap number.
    //        Best to run them now.

    // Clear the old mapping - but keep fte->page (the SOS window cap)
    bzero(&(data->fte->mapping), sizeof(struct frame_mapping));
    // assert(!data->fte->mapping.proc_cap);

    dprintf(6, "frame_swapout_finalise: swapped %dB to swap %d from frame %d\n",
            data->sent, data->swap_num, _fte_to_framenumber(data->fte));

    // No longer in flight
    data->meta->inflight = 0;

    // Trigger original callback. It gets the freed frame.
    (*(data->cb))(FRAME_TABLE_SUCCESS, _fte_to_framenumber(data->fte),
                  data->swap_num, data->cookie);

    // Free heap storage
    free(data);

    return SOS_TASK_SUCCESS;
}

// This is the public interface for swapping out frames. It searches the
// swap queue and sets up asynchronous ops where necessary.
// Swapout should be called from frame_alloc when the free list is empty
int frame_swapout(frame_swapout_cb cb, void *cookie) {
    frame_table_entry *tmp;
    frame_swapout_data *data;
    int count = 0;

    dprintf(5, "frame_swapout: looking for a frame to evict...\n");

    // Malloc some heap space to finalise this operation
    data = malloc(sizeof(struct frame_swapout_data));
    if(data == NULL) {
        // XXX Nothing can continue from here
        dprintf(5, "frame_swapout: out of memory, cannot proceed\n");
        return FRAME_TABLE_OOM;
    }
    // Copy useful stuff to the heap
    data->cb = cb;
    data->cookie = cookie;
    data->sent = 0;

    tmp = _swap_queue;
    while (true) {
        if (tmp && tmp->referenced) {
            // First chance. Unmap it from all processes,
            _frame_chance(tmp);
            // ... then put at the back of the queue.
            DL_DELETE(_swap_queue, tmp);
            tmp->next = NULL;
            tmp->prev = NULL;
            assert(tmp->swappable);
            DL_APPEND(_swap_queue, tmp);
        } else if (tmp) {
            // Second chance - pop from the queue.
            dprintf(7, "frame_swapout: 2nd chance for %p\n", tmp);
            DL_DELETE(_swap_queue, tmp);
            tmp->next = NULL;
            tmp->prev = NULL;
            assert(tmp->swappable);
            tmp->swappable = 0;
            // swap out tmp. frame_swapout will perform the client callback.
            data->fte = tmp;
            _frame_swapout(data);
            break;
        } else if (_swap_queue) {
            // At least one ref'd frame is avail to swap. Take the queue head.
            dprintf(7, "frame_swapout: Choosing queue head\n");
            tmp = _swap_queue;
            data->fte = tmp;
            DL_DELETE(_swap_queue, tmp);
            tmp->next = NULL;
            tmp->prev = NULL;
            assert(tmp->swappable);
            tmp->swappable = 0;
            // This frame is almost certainly ref'dd and needs to be unmapped.
            if (tmp->referenced) {
                _frame_chance(tmp);
            }
            _frame_swapout(data);
            break;
        } else {
            // swap queue is completely empty. very bad.
            dprintf(0, "SOS is out of swappable frames.\n");
            return FRAME_TABLE_OOM;
        }
        tmp = tmp->next;
        count++;
    }

    dprintf(6, "frame_swapout: considered %d in the queue\n", count);
    dprintf(6, "frame_swapout: %d frames managed\n", _cur_frames);

    return FRAME_TABLE_SUCCESS;
}

/* Delete metadata and forget about a frame that's in swap file */
int free_swapped_frame(seL4_Word swap_number) {
    int err;
    swapped_frames *meta = _swap_meta + swap_number;
    bzero((void *) meta, sizeof(struct swapped_frames));
    return FRAME_TABLE_SUCCESS;
}

/* Public interface to restore a frame from swap */
int frame_swapin(seL4_Word swap_number, frame_swapin_cb cb, void *cookie) {
    int err;
    swapped_frames *meta = _swap_meta + swap_number;
    frame_swapin_data *data;

    dprintf(5, "frame_swapin: SN %d\n", swap_number);

    // Check if the swapout is still in progress, and abort if so
    if (meta->inflight) {
        meta->abort = 1;
        meta->cb = cb;
        meta->cookie = cookie;
        return FRAME_TABLE_SUCCESS;
    }

    // Grab space on the heap
    data = malloc(sizeof(struct frame_swapin_data));
    if (data == NULL) {
        // Out of kernel memory - oops.
        dprintf(1, "frame_swapin: out of kernel memory\n");
        return FRAME_TABLE_INVAL;
    }
    bzero((void *) data, sizeof(struct frame_swapin_data));

    // Put useful stuff into the heap
    data->sent = 0;
    data->cb = cb;
    data->cookie = cookie;
    data->swap_number = swap_number;
    data->meta = meta;

    // Get a frame with which to proceed
    err = frame_alloc(&_swapin_gotframe, (void *) data);

    return err;
}

static void _swapin_gotframe(seL4_Word sos_vaddr, void *cookie) {
    int err;
    frame_swapin_data *data = (frame_swapin_data *) cookie;

    dprintf(6, "frame_swapin: got fresh frame at %p\n", sos_vaddr);

    // Put the frame number into data
    data->sos_vaddr = sos_vaddr;
    data->frame_number = frame_translate_vaddr(sos_vaddr);

    // Set up first NFS read call
    err = nfs_read(&_swap_file, data->swap_number * PAGE_SIZE, PAGE_SIZE,
                   &_finalise_swapin, (uintptr_t) cookie);
    conditional_panic(err, "frame_swapin: nfs_read failed\n");
    return;
}

/* Callback passed to NFS Read for reading a frame */
static void _finalise_swapin(uintptr_t token, enum nfs_stat status,
                             fattr_t *fattr, int count, void *frame) {
    int err = SOS_PAGE_TABLE_SUCCESS;
    frame_swapin_data *data = (frame_swapin_data *) token;

    dprintf(7, "frame_swapin_finalise: got %dB\n", count);

    // Check for abort
    if (data->meta->abort || data->meta->mapping.pcb->abort) {
        dprintf(6, "frame_swapin_finalise: aborting\n");
        (*(data->cb))(SOS_PAGE_TABLE_SEGFAULT, data->cookie);
        free(data);
        return;
    }

    // Check status
    if (status) {
        dprintf(1, "frame_swapin_finalise: error %d\n", status);
        free(data);
        return;
    }

    // Copy data back into our frame
    memcpy((void *) (data->sos_vaddr + data->sent), frame, count);
    data->sent += count;

    // Update count
    if (data->sent < PAGE_SIZE) {
        // Set up next call if it's necessary
        err = (int) nfs_read(&_swap_file,
                             (data->swap_number * PAGE_SIZE) + data->sent,
                             PAGE_SIZE - data->sent,
                             &_finalise_swapin, (uintptr_t) token);
        if (err) {
            dprintf(1, "frame_swapin_finalise: couldn't set up further read\n");
            return;
        }
    } else {

        dprintf(6, "frame_swapin_finalise: swap %d is back at %p\n",
                data->swap_number, data->sos_vaddr);

        // Call _frame_swapin to update the details
        err = _frame_swapin(data->swap_number, data->sos_vaddr);

        // Perform callback
        (*(data->cb))(err, data->cookie);

        // Free heap data
        free(data);

        return;
    }
}

/* Update mappings, metadata and page tables after a frame retrieved. */
/* Called once the frame has been fully written to sos_vaddr */
static int _frame_swapin(seL4_Word swap_num, seL4_Word sos_vaddr) {
    int err;
    swapped_frames *meta = _swap_meta + swap_num;
    frame_table_entry *fte = _vaddr_to_fte(sos_vaddr);

    // Return swap_num to the swap table
    sos_swap_free(swap_num);

    // re-do the process grant
    err = frame_grant(sos_vaddr, meta->mapping.pcb, meta->mapping.proc_vaddr);
    if (err != FRAME_TABLE_SUCCESS) {
        // Something truly terrible has happened
        dprintf(1, "frame_restore: Could not grant to process\n");
        return FRAME_TABLE_INVAL;
    }

    // Perform a page_swapin - updates page table with new info
    err = sos_page_swapin(meta->mapping.pcb, meta->mapping.proc_vaddr,
                          frame_translate_vaddr(sos_vaddr));
    seL4_CPtr proc_cap;
    frame_process_cap(sos_vaddr, meta->mapping.pcb, &proc_cap);
    err = sos_page_map(meta->mapping.pcb, meta->mapping.proc_vaddr, proc_cap);

    // Flush icache again for some reason (frame_alloc already flushed it once)
    frame_flush_icache(frame_translate_vaddr(sos_vaddr));

    if (meta->abort) {
        dprintf(6, "frame_swapin: Waking up faulted process after abort\n");
        // We perform an extra callback for the faulted (waiting) process
        (*(meta->cb))(0, meta->cookie);
    }

    return FRAME_TABLE_SUCCESS;
}
